/* Copyright 2012 Tobias Marschall
 *
 * This file is part of CLEVER.
 *
 * CLEVER is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * CLEVER is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with CLEVER.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <sstream>

#include "GroupWiseBamReader.h"

using namespace std;

GroupWiseBamReader::GroupWiseBamReader(const std::string& filename, bool paired) : read_name("") {
	this->counter = 0;
	this->progress_messages_os = 0;
	this->progress_message_frequency = -1;
	this->skipped_duplicates = 0;
	this->paired = paired;
	this->unmapped1 = true;
	this->unmapped2 = true;
	this->inconsistentXTags = false;
	this->inconsistentXTagsMessage = "";
	this->multiplyMappedButNoXA = false;
	if (!bam_reader.Open(filename)) {
		ostringstream oss;
		oss << "Could not read BAM input from \"" << filename << "\".";
		throw std::runtime_error(oss.str());
	}
	this->finished = !bam_reader.GetNextAlignment(next_read_aln);
}

GroupWiseBamReader::~GroupWiseBamReader() {
}


long long GroupWiseBamReader::hashReadStart(const BamTools::BamAlignment& aln) {
	assert(sizeof(long long) >= 2*sizeof(int32_t));
	long long result = aln.Position;
	result |= ((long long)aln.RefID) << (8*sizeof(int32_t));
	return result;
}

void GroupWiseBamReader::advance() {
	assert(!finished);
	read_name = next_read_aln.Name;
	alignments1.clear();
	alignments2.clear();
	// start positions of the current alignments (stored to check for
	// duplicate alignments)
	boost::unordered_set<long long> alignments1_starts;
	boost::unordered_set<long long> alignments2_starts;
	// number of alignments (including those skipped because they are duplicates)
	int alignments1_count = 0;
	int alignments2_count = 0;
	int n = 0;
	this->unmapped1 = true;
	this->unmapped2 = true;
	this->inconsistentXTags = false;
	this->inconsistentXTagsMessage = "";
	this->multiplyMappedButNoXA = false;
	while (true) {
		if (n++ > 0) {
			if (!bam_reader.GetNextAlignment(next_read_aln)){
				finished = true;
				break;
			}
			if (read_name.compare(next_read_aln.Name) != 0) {
				if (read_names.find(next_read_aln.Name) != read_names.end()) {
					ostringstream oss;
					oss << "Error: Reads not grouped properly. Offending read: \"" << next_read_aln.Name << "\"" << endl;
					throw std::runtime_error(oss.str());
				}
				// store the first million read names to check whether reads are grouped properly.
				if (read_names.size() < 1000000) {
					read_names.insert(next_read_aln.Name);
				}
				break;
			}
		}
		counter += 1;
		if ((progress_messages_os != 0) && (counter % progress_message_frequency == 0)) {
			*progress_messages_os << "Having processed " << counter << " read alignments" << endl;
		}
		if (paired) {
			assert(next_read_aln.IsFirstMate() != next_read_aln.IsSecondMate());
		}
		uint32_t x0 = 0;
		uint32_t x1 = 0;
		if (next_read_aln.GetTag("X0", x0) && next_read_aln.GetTag("X1", x1)) {
			string xa = "";
			if (!next_read_aln.GetTag("XA", xa) && (x0+x1>1)) {
				multiplyMappedButNoXA = true;
			}
		}
		if (!paired || next_read_aln.IsFirstMate()) {
			if (next_read_aln.IsMapped()) unmapped1 = false;
			addRead(alignments1, alignments1_starts, next_read_aln);
			alignments1_count += 1;
		} else {
			if (next_read_aln.IsMapped()) unmapped2 = false;
			addRead(alignments2, alignments2_starts, next_read_aln);
			alignments2_count += 1;
		}
	}
	checkX0X1Tags(alignments1,alignments1_count);
	if (paired) {
		checkX0X1Tags(alignments2,alignments2_count);
	}
}

bool GroupWiseBamReader::hasNext() const {
	return !finished;
}

const std::vector<BamTools::BamAlignment>& GroupWiseBamReader::getAlignments() const {
	assert(!paired);
	return alignments1;
}

const vector<BamTools::BamAlignment>& GroupWiseBamReader::getAlignmentsFirst() const {
	assert(paired);
	return alignments1;
}

const vector<BamTools::BamAlignment>& GroupWiseBamReader::getAlignmentsSecond() const {
	assert(paired);
	return alignments2;
}

bool GroupWiseBamReader::isUnmapped() const {
	assert(!paired);
	return unmapped1;
}

bool GroupWiseBamReader::isFirstUnmapped() const {
	assert(paired);
	return unmapped1;
}

bool GroupWiseBamReader::isSecondUnmapped() const {
	assert(paired);
	return unmapped2;
}

bool GroupWiseBamReader::hasInconsistentXTags() const {
	return inconsistentXTags;
}

bool GroupWiseBamReader::hasMultipleMappingsButNoXA() const {
	return multiplyMappedButNoXA;
}

const std::string& GroupWiseBamReader::getInconsistentXTagsMessage() const {
	assert(inconsistentXTags);
	return inconsistentXTagsMessage;
}


void GroupWiseBamReader::enableProgressMessages(std::ostream& os, int frequency) {
	assert(frequency > 0);
	this->progress_messages_os = &os;
	this->progress_message_frequency = frequency;
}

long long GroupWiseBamReader::getSkippedDuplicates() const {
	return skipped_duplicates;
}

void GroupWiseBamReader::addRead(vector<BamTools::BamAlignment>& alignments, boost::unordered_set<long long>& alignments_starts, const BamTools::BamAlignment& read_aln) {
	long long start_hash = hashReadStart(read_aln);
	if (alignments_starts.find(start_hash) != alignments_starts.end()) {
		skipped_duplicates += 1;
		// cerr << "Skipping "  << read_aln.Name << endl;
	} else {
		alignments.push_back(read_aln);
		alignments_starts.insert(start_hash);
	}
}

/** Verifies whether the number of reads complies with the given X0/X1 tags (if given). */
void GroupWiseBamReader::checkX0X1Tags(const vector<BamTools::BamAlignment>& alignments, int count) {
	vector<BamTools::BamAlignment>::const_iterator it = alignments.begin();
	for (;it!=alignments.end();++it) {
		uint32_t x0 = 0;
		uint32_t x1 = 0;
		if (it->GetTag("X0", x0) && it->GetTag("X1", x1)) {
			if (count != (int)(x0+x1)) {
				ostringstream oss;
				oss << "Alignment count mismatch for read \"" << it->Name << "\": ";
				oss << "X0=" << x0 << ", X1=" << x1 << ", but found " << count << " alignment(s).";
				inconsistentXTagsMessage = oss.str();
				inconsistentXTags = true;
				return;
			}
		}
	}
}
